"""
数据库连接和会话管理
"""

from typing import AsyncGenerator, Optional, Dict, Any
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import StaticPool
from sqlalchemy import text
from contextlib import asynccontextmanager
import asyncio
import logging
from contextvars import ContextVar

from ..core.config import settings
from ..core.logging import get_logger
from ..core.exceptions import DatabaseError

logger = get_logger(__name__)

# 上下文变量用于存储当前会话
session_context: ContextVar[Optional[AsyncSession]] = ContextVar('session', default=None)


class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self):
        self.engine = None
        self.session_factory = None
        self._initialized = False
        self._lock = asyncio.Lock()
    
    async def initialize(self, database_url: Optional[str] = None):
        """初始化数据库连接"""
        async with self._lock:
            if self._initialized:
                return
            
            try:
                db_url = database_url or settings.DATABASE_URL
                
                # 创建数据库引擎
                if db_url.startswith("sqlite"):
                    # SQLite特殊配置
                    self.engine = create_async_engine(
                        db_url,
                        echo=settings.DATABASE_ECHO,
                        poolclass=StaticPool,
                        connect_args={
                            "check_same_thread": False,
                            "timeout": 30
                        }
                    )
                else:
                    # 其他数据库配置
                    self.engine = create_async_engine(
                        db_url,
                        echo=settings.DATABASE_ECHO,
                        pool_pre_ping=True,
                        pool_recycle=3600,
                        pool_size=10,
                        max_overflow=20
                    )
                
                # 创建会话工厂
                self.session_factory = async_sessionmaker(
                    bind=self.engine,
                    class_=AsyncSession,
                    expire_on_commit=False,
                    autoflush=True
                )
                
                self._initialized = True
                logger.info("Database manager initialized successfully")
                
            except Exception as e:
                logger.error(f"Failed to initialize database manager: {e}")
                raise DatabaseError(f"Database initialization failed: {str(e)}")
    
    async def create_session(self) -> AsyncSession:
        """创建数据库会话"""
        if not self._initialized:
            await self.initialize()
        
        if not self.session_factory:
            raise DatabaseError("Session factory not initialized")
        
        return self.session_factory()
    
    @asynccontextmanager
    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        """获取数据库会话上下文管理器"""
        session = None
        try:
            session = await self.create_session()
            session_context.set(session)
            yield session
            await session.commit()
        except Exception as e:
            if session:
                await session.rollback()
            logger.error(f"Database session error: {e}")
            raise
        finally:
            if session:
                await session.close()
                session_context.set(None)
    
    def get_current_session(self) -> Optional[AsyncSession]:
        """获取当前会话"""
        return session_context.get()
    
    async def close(self):
        """关闭数据库连接"""
        if self.engine:
            await self.engine.dispose()
            self._initialized = False
            logger.info("Database manager closed")
    
    async def health_check(self) -> Dict[str, Any]:
        """数据库健康检查"""
        try:
            async with self.get_session() as session:
                result = await session.execute(text("SELECT 1"))
                await result.fetchone()
                
            return {
                "status": "healthy",
                "message": "Database connection is working",
                "url": str(self.engine.url).replace(self.engine.url.password or "", "***") if self.engine else None
            }
        except Exception as e:
            return {
                "status": "unhealthy",
                "message": f"Database connection failed: {str(e)}"
            }
    
    async def get_connection_info(self) -> Dict[str, Any]:
        """获取连接信息"""
        if not self.engine:
            return {"status": "not_initialized"}
        
        pool = self.engine.pool
        
        return {
            "status": "initialized",
            "url": str(self.engine.url).replace(self.engine.url.password or "", "***"),
            "driver": self.engine.driver,
            "dialect": self.engine.dialect.name,
            "pool_size": getattr(pool, "size", 0),
            "checked_in": getattr(pool, "checkedin", 0),
            "checked_out": getattr(pool, "checkedout", 0),
            "overflow": getattr(pool, "overflow", 0)
        }


# 全局数据库管理器实例
db_manager = DatabaseManager()


# 依赖注入函数
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI依赖注入函数"""
    async for session in db_manager.get_session():
        yield session


# 事务装饰器
def transactional(func):
    """事务装饰器"""
    async def wrapper(*args, **kwargs):
        async with db_manager.get_session() as session:
            # 如果函数有session参数，注入当前会话
            if 'session' in func.__code__.co_varnames:
                kwargs['session'] = session
            return await func(*args, **kwargs)
    
    return wrapper


def get_current_db_session() -> Optional[AsyncSession]:
    """获取当前数据库会话"""
    return db_manager.get_current_session()


# 数据库操作基类
class DatabaseOperation:
    """数据库操作基类"""
    
    def __init__(self, session: Optional[AsyncSession] = None):
        self.session = session or get_current_db_session()
        
        if not self.session:
            raise DatabaseError("No database session available")
    
    async def execute(self, query, params: Optional[Dict[str, Any]] = None):
        """执行查询"""
        try:
            if params:
                result = await self.session.execute(query, params)
            else:
                result = await self.session.execute(query)
            return result
        except Exception as e:
            logger.error(f"Database execute error: {e}")
            raise DatabaseError(f"Query execution failed: {str(e)}")
    
    async def fetch_one(self, query, params: Optional[Dict[str, Any]] = None):
        """获取单条记录"""
        result = await self.execute(query, params)
        return result.fetchone()
    
    async def fetch_all(self, query, params: Optional[Dict[str, Any]] = None):
        """获取所有记录"""
        result = await self.execute(query, params)
        return result.fetchall()
    
    async def fetch_scalar(self, query, params: Optional[Dict[str, Any]] = None):
        """获取标量值"""
        result = await self.execute(query, params)
        return result.scalar()
    
    async def insert(self, model_instance):
        """插入记录"""
        try:
            self.session.add(model_instance)
            await self.session.flush()  # 获取ID但不提交
            return model_instance
        except Exception as e:
            logger.error(f"Database insert error: {e}")
            raise DatabaseError(f"Insert operation failed: {str(e)}")
    
    async def update(self, model_instance):
        """更新记录"""
        try:
            self.session.merge(model_instance)
            await self.session.flush()
            return model_instance
        except Exception as e:
            logger.error(f"Database update error: {e}")
            raise DatabaseError(f"Update operation failed: {str(e)}")
    
    async def delete(self, model_instance):
        """删除记录"""
        try:
            await self.session.delete(model_instance)
            await self.session.flush()
        except Exception as e:
            logger.error(f"Database delete error: {e}")
            raise DatabaseError(f"Delete operation failed: {str(e)}")
    
    async def commit(self):
        """提交事务"""
        try:
            await self.session.commit()
        except Exception as e:
            await self.session.rollback()
            logger.error(f"Database commit error: {e}")
            raise DatabaseError(f"Commit operation failed: {str(e)}")
    
    async def rollback(self):
        """回滚事务"""
        try:
            await self.session.rollback()
        except Exception as e:
            logger.error(f"Database rollback error: {e}")


# 连接池管理
class ConnectionPoolManager:
    """连接池管理器"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db_manager = db_manager
        self._monitoring = False
    
    async def start_monitoring(self, interval: int = 60):
        """开始监控连接池"""
        self._monitoring = True
        
        while self._monitoring:
            try:
                info = await self.db_manager.get_connection_info()
                logger.info(f"Connection pool status: {info}")
                await asyncio.sleep(interval)
            except Exception as e:
                logger.error(f"Connection pool monitoring error: {e}")
                await asyncio.sleep(interval)
    
    def stop_monitoring(self):
        """停止监控连接池"""
        self._monitoring = False
    
    async def get_pool_stats(self) -> Dict[str, Any]:
        """获取连接池统计"""
        if not self.db_manager.engine:
            return {"status": "not_initialized"}
        
        pool = self.db_manager.engine.pool
        
        return {
            "size": getattr(pool, "size", 0),
            "checked_in": getattr(pool, "checkedin", 0),
            "checked_out": getattr(pool, "checkedout", 0),
            "overflow": getattr(pool, "overflow", 0),
            "invalid": getattr(pool, "invalid", 0),
            "max_overflow": getattr(pool, "_max_overflow", 0)
        }


# 全局连接池管理器
pool_manager = ConnectionPoolManager(db_manager)


# 初始化函数
async def init_database():
    """初始化数据库"""
    try:
        await db_manager.initialize()
        
        # 开始连接池监控
        asyncio.create_task(pool_manager.start_monitoring())
        
        logger.info("Database initialized successfully")
        return True
        
    except Exception as e:
        logger.error(f"Failed to initialize database: {e}")
        return False


# 清理函数
async def cleanup_database():
    """清理数据库资源"""
    try:
        pool_manager.stop_monitoring()
        await db_manager.close()
        logger.info("Database cleanup completed")
    except Exception as e:
        logger.error(f"Database cleanup error: {e}")
